package org.framework.lazy.cloud.network.heartbeat.protocol;

import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.DatagramPacket;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioDatagramChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.http.*;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;

import java.net.InetSocketAddress;
import java.net.URI;

public class MultiProtocolProxyServer {
    private static final int HTTP_PORT = 8080;
    private static final int TCP_PORT = 9090;
    private static final int UDP_PORT = 10000;

    public static void main(String[] args) throws Exception {
        EventLoopGroup bossGroup = new NioEventLoopGroup();
        EventLoopGroup workerGroup = new NioEventLoopGroup();

        try {
            // 启动 HTTP 代理服务器
            startHttpProxy(bossGroup, workerGroup);
            // 启动 TCP 代理服务器
            startTcpProxy(bossGroup, workerGroup);
            // 启动 UDP 代理服务器
            startUdpProxy(workerGroup);

            System.out.println("Proxy servers started.");
        } finally {
            bossGroup.shutdownGracefully();
            workerGroup.shutdownGracefully();
        }
    }

    private static void startHttpProxy(EventLoopGroup bossGroup, EventLoopGroup workerGroup) throws InterruptedException {
        ServerBootstrap b = new ServerBootstrap();
        b.group(bossGroup, workerGroup)
                .channel(NioServerSocketChannel.class)
                .handler(new LoggingHandler(LogLevel.INFO))
                .childHandler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) throws Exception {
                        ch.pipeline().addLast(
                                new HttpClientCodec(),
                                new HttpObjectAggregator(1048576),
                                new HttpProxyServerHandler()
                        );
                    }
                });

        b.bind(HTTP_PORT).sync();
        System.out.println("HTTP proxy server started on port " + HTTP_PORT);
    }

    private static void startTcpProxy(EventLoopGroup bossGroup, EventLoopGroup workerGroup) throws InterruptedException {
        ServerBootstrap b = new ServerBootstrap();
        b.group(bossGroup, workerGroup)
                .channel(NioServerSocketChannel.class)
                .handler(new LoggingHandler(LogLevel.INFO))
                .childHandler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) throws Exception {
                        ch.pipeline().addLast(new TcpProxyServerHandler());
                    }
                });

        b.bind(TCP_PORT).sync();
        System.out.println("TCP proxy server started on port " + TCP_PORT);
    }

    private static void startUdpProxy(EventLoopGroup workerGroup) throws InterruptedException {
        Bootstrap b = new Bootstrap();
        b.group(workerGroup)
                .channel(NioDatagramChannel.class)
                .handler(new UdpProxyServerHandler());

        b.bind(UDP_PORT).sync();
        System.out.println("UDP proxy server started on port " + UDP_PORT);
    }

    private static class HttpProxyServerHandler extends ChannelInboundHandlerAdapter {
        private Channel outboundChannel;

        @Override
        public void channelRead(final ChannelHandlerContext ctx, Object msg) throws Exception {
            if (msg instanceof FullHttpRequest) {
                FullHttpRequest request = (FullHttpRequest) msg;
                URI uri = new URI(request.uri());
                String host = uri.getHost();
                int port = uri.getPort();
                if (port == -1) {
                    port = 80;
                }

                Bootstrap b = new Bootstrap();
                b.group(ctx.channel().eventLoop())
                        .channel(NioSocketChannel.class)
                        .handler(new ChannelInitializer<SocketChannel>() {
                            @Override
                            protected void initChannel(SocketChannel ch) throws Exception {
                                ch.pipeline().addLast(
                                        new HttpClientCodec(),
                                        new HttpObjectAggregator(1048576),
                                        new HttpProxyBackendHandler(ctx.channel())
                                );
                            }
                        });

                ChannelFuture f = b.connect(host, port);
                outboundChannel = f.channel();
                f.addListener(new ChannelFutureListener() {
                    @Override
                    public void operationComplete(ChannelFuture future) throws Exception {
                        if (future.isSuccess()) {
                            outboundChannel.writeAndFlush(request.retain());
                        } else {
                            ctx.channel().close();
                        }
                    }
                });
            }
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            if (outboundChannel != null) {
                closeOnFlush(outboundChannel);
            }
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            cause.printStackTrace();
            closeOnFlush(ctx.channel());
        }

        static void closeOnFlush(Channel ch) {
            if (ch.isActive()) {
                ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            }
        }
    }

    private static class HttpProxyBackendHandler extends ChannelInboundHandlerAdapter {
        private final Channel inboundChannel;

        HttpProxyBackendHandler(Channel inboundChannel) {
            this.inboundChannel = inboundChannel;
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
            inboundChannel.writeAndFlush(msg);
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            closeOnFlush(inboundChannel);
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            cause.printStackTrace();
            closeOnFlush(ctx.channel());
        }

        static void closeOnFlush(Channel ch) {
            if (ch.isActive()) {
                ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            }
        }
    }

    private static class TcpProxyServerHandler extends ChannelInboundHandlerAdapter {
        private Channel outboundChannel;

        @Override
        public void channelRead(final ChannelHandlerContext ctx, Object msg) throws Exception {
            if (outboundChannel == null) {
                // 假设目标服务器地址和端口
                String targetHost = "example.com";
                int targetPort = 80;

                Bootstrap b = new Bootstrap();
                b.group(ctx.channel().eventLoop())
                        .channel(NioSocketChannel.class)
                        .handler(new ChannelInitializer<SocketChannel>() {
                            @Override
                            protected void initChannel(SocketChannel ch) throws Exception {
                                ch.pipeline().addLast(new TcpProxyBackendHandler(ctx.channel()));
                            }
                        });

                ChannelFuture f = b.connect(targetHost, targetPort);
                outboundChannel = f.channel();
                f.addListener(new ChannelFutureListener() {
                    @Override
                    public void operationComplete(ChannelFuture future) throws Exception {
                        if (future.isSuccess()) {
                            outboundChannel.writeAndFlush(msg);
                        } else {
                            ctx.channel().close();
                        }
                    }
                });
            } else {
                outboundChannel.writeAndFlush(msg);
            }
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            if (outboundChannel != null) {
                closeOnFlush(outboundChannel);
            }
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            cause.printStackTrace();
            closeOnFlush(ctx.channel());
        }

        static void closeOnFlush(Channel ch) {
            if (ch.isActive()) {
                ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            }
        }
    }

    private static class TcpProxyBackendHandler extends ChannelInboundHandlerAdapter {
        private final Channel inboundChannel;

        TcpProxyBackendHandler(Channel inboundChannel) {
            this.inboundChannel = inboundChannel;
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
            inboundChannel.writeAndFlush(msg);
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            closeOnFlush(inboundChannel);
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            cause.printStackTrace();
            closeOnFlush(ctx.channel());
        }

        static void closeOnFlush(Channel ch) {
            if (ch.isActive()) {
                ch.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
            }
        }
    }

    private static class UdpProxyServerHandler extends SimpleChannelInboundHandler<DatagramPacket> {
        private InetSocketAddress targetAddress = new InetSocketAddress("example.com", 80);

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, DatagramPacket packet) throws Exception {
            Bootstrap b = new Bootstrap();
            b.group(ctx.channel().eventLoop())
                    .channel(NioDatagramChannel.class)
                    .handler(new UdpProxyBackendHandler(ctx.channel(), packet.sender()));

            Channel channel = b.bind(0).sync().channel();
            channel.writeAndFlush(new DatagramPacket(packet.content().retain(), targetAddress));
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            cause.printStackTrace();
            ctx.close();
        }
    }

    private static class UdpProxyBackendHandler extends SimpleChannelInboundHandler<DatagramPacket> {
        private final Channel inboundChannel;
        private final InetSocketAddress clientAddress;

        UdpProxyBackendHandler(Channel inboundChannel, InetSocketAddress clientAddress) {
            this.inboundChannel = inboundChannel;
            this.clientAddress = clientAddress;
        }

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, DatagramPacket packet) throws Exception {
            inboundChannel.writeAndFlush(new DatagramPacket(packet.content().retain(), clientAddress));
            ctx.channel().close();
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
            cause.printStackTrace();
            ctx.close();
        }
    }
}